import dgl
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

"""
数据的预处理，去掉偏差比较大的数据
"""

def del_errors(df, max_error=5):  # 默认类型是pd
    ids = []  # 记录误差较大数据的下标
    for i in range(len(df)):
        sample = df.loc[i]
        reactivity_error = sample['reactivity_error']
        deg_error_Mg_pH10 = sample['deg_error_Mg_pH10']
        deg_error_Mg_50C = sample['deg_error_Mg_50C']
        reactivity_error_mean = np.mean(reactivity_error)
        deg_error_Mg_pH10_mean = np.mean(deg_error_Mg_pH10)
        deg_error_Mg_50C_mean = np.mean(deg_error_Mg_50C)
        if reactivity_error_mean >= max_error or \
                deg_error_Mg_pH10_mean >= max_error or \
                deg_error_Mg_50C_mean >= max_error:
            ids.append(i)
    # print(ids)
    # print(len(ids))
    df = df.drop(labels=ids, axis=0)
    df = df.reset_index()
    return df

def plat_errors(df):
    """画图展示偏差"""
    means = {0:[], 1:[], 2:[]}
    for i in range(len(df)):
        sample = df.loc[i]
        reactivity_error = sample['reactivity_error']
        deg_error_Mg_pH10 = sample['deg_error_Mg_pH10']
        deg_error_Mg_50C = sample['deg_error_Mg_50C']
        reactivity_error_mean = np.mean(reactivity_error)
        deg_error_Mg_pH10_mean = np.mean(deg_error_Mg_pH10)
        deg_error_Mg_50C_mean = np.mean(deg_error_Mg_50C)
        means[0].append(reactivity_error_mean)
        means[1].append(deg_error_Mg_pH10_mean)
        means[2].append(deg_error_Mg_50C_mean)
    colors = ['blue', 'red', 'green']
    y_labels = ['reactivity_error', 'deg_Mg_pH10_error', 'deg_Mg_50C_error']
    for i in range(3):
        plt.figure(i)
        x = range(len(df))
        y = means[i]
        plt.plot(x, y, label='error', color=colors[i])
        plt.xlabel('index')
        plt.ylabel(y_labels[i])
        plt.legend(loc='best')
    plt.show()

if __name__ == '__main__':
    # 读取数据集, 调试和检查
    data_path = "../dataset/valid.json"
    df = pd.read_json(data_path, lines=True)
    sample = df.loc[0]

    reactivity_error_mean = []
    reactivity_error_list = []
    for i in range(len(df)):
        sample = df.loc[i]
        reactivity_error = sample['reactivity_error']
        reactivity_error_list.append(reactivity_error)
        reactivity_error_mean.append(np.mean(reactivity_error))

    reactivity_error_list = np.array(reactivity_error_list)
    reactivity_error_mean = np.array(reactivity_error_mean)
    print(len(df))
    plat_errors(df)
    df = del_errors(df, max_error=5)
    print(len(df))
    plat_errors(df)



